#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat May  8 03:08:38 2021

@author: sarupa
"""

#Imports 
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from LSTM_model import *
from sklearn.metrics import mean_squared_error
from casadi import *
from numpy import linalg as la
import control
import timeit
import time
from numpy import random
from random import seed

start = timeit.default_timer()


def rmse(predictions, targets):

    differences = (predictions - targets)/targets                      

    differences_squared = differences ** 2                    

    mean_of_differences_squared = differences_squared.mean()  

    rmse_val = np.sqrt(mean_of_differences_squared)*100         

    return rmse_val 


start1 = time.perf_counter()
# Load the weights, units and LSTM properties
W_1 = np.loadtxt('./Saved Models3/W_1.txt')
W_2 = np.loadtxt('./Saved Models3/W_2.txt')
U_1 = np.loadtxt('./Saved Models3/U_1.txt')
U_2 = np.loadtxt('./Saved Models3/U_2.txt')
b_1 = np.loadtxt('./Saved Models3/b_1.txt')
b_2 = np.loadtxt('./Saved Models3/b_2.txt')
V_m = np.loadtxt('./Saved Models3/V_m.txt')
b_m = np.loadtxt('./Saved Models3/b_m.txt')

units_1=50
units_2=50
sequenceLength=2
numberOfFeatures = 10
numberOfOutputs = 1
sequenceLength_out=2
number_of_states= 4

# Load the testing data
y_test = np.loadtxt("./Saved Models3/y_test.txt")
y_test_unscaled = np.loadtxt("./Saved Models3/y_test_unscaled.txt")
X_test = np.loadtxt("./Saved Models3/X_test.txt")
X_test = X_test.reshape(X_test.shape[0], sequenceLength, numberOfFeatures)
data_max=np.loadtxt('./Saved Models3/data_max.txt')
data_min=np.loadtxt('./Saved Models3/data_min.txt')

predicted_trajectory_unscaled=np.loadtxt('./Saved Models3/predicted_trajectory_unscaled.txt')
y_test_MSAP_unscaled=np.loadtxt('./Saved Models3/y_test_MSAP_unscaled.txt')

#from random import random
np.random.seed(100) 



#Define array
Nsim = y_test.shape[0]-sequenceLength
u_model = X_test[0,:,number_of_states:numberOfFeatures] # input to the model

# Define std deviation
sigma_v = 20 # Standard deviation of the measurements
sigma_w = 0.005 # Standard deviation for the process noise
sigma_p = 100 # Standard deviation for prior

Q = np.diag((sigma_w*np.ones((number_of_states ,)))**2)
Rc = np.diag((sigma_v*np.ones((2,)))**2)

sigma_v_m=0.005
v = sigma_v_m*random.randn(Nsim,2,sequenceLength)

Pc = np.ones((number_of_states,number_of_states,Nsim+1))
Pc[:,:,0]=np.diag((sigma_p)*np.ones((number_of_states ,)))
Pc[:,:,1]=np.diag((sigma_p)*np.ones((number_of_states ,)))

# Define Casadi variables
X_hat1=SX.sym('X',(1,number_of_states))
X_hat2=SX.sym('X',(1,number_of_states))
u_hat=SX.sym('U',(sequenceLength,numberOfFeatures-number_of_states))

# Define symbolic function
jac_sym=ObtainlstmModel_casadi_1(X_hat1,X_hat2,u_hat,W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength)

# Find jacobians
jac_fun1=jacobian(jac_sym,X_hat1)
Ad_fun1 = Function('Ad_fun1',[X_hat1,X_hat2,u_hat],[jac_fun1])

jac_fun2=jacobian(jac_sym,X_hat2)
Ad_fun2 = Function('Ad_fun2',[X_hat1,X_hat2,u_hat],[jac_fun2])

jac_fun3=jacobian(jac_sym,u_hat)
Bd_fun = Function('Bd_fun',[X_hat1,X_hat2,u_hat],[jac_fun3])

# Deifne initial guess for ekf and open loop simulation for comparison
x_model=X_test[0,:,0:number_of_states]*0.95
x_model_op=X_test[0,:,0:number_of_states]*0.95

# Data array
x_est=np.zeros((number_of_states,Nsim))
x_opl=np.zeros((number_of_states,Nsim))
x_est[:,0:sequenceLength]= x_model.T
x_opl[:,0:sequenceLength]= x_model_op.T

# Actual data save
X_act=np.zeros((number_of_states,Nsim))
X_act[:,0:sequenceLength] =X_test[0,:,0:number_of_states].T

# Evaluation of Jacobians at the initial guess
Ad1= np.array(Ad_fun1(x_model[0],x_model[1],u_model))
Ad2= np.array(Ad_fun2(x_model[0],x_model[1],u_model))
Bd= np.array(Bd_fun(x_model[0],x_model[1],u_model))

Pc[:,:,2] = Ad1.dot(Pc[:,:,0]).dot(np.transpose(Ad1))+ Ad2.dot(Pc[:,:,1]).dot(np.transpose(Ad2))+ Q #+Ad1.dot(Pc[:,:,1]).dot(np.transpose(Ad2))+ Ad2.dot(Pc[:,:,0]).dot(np.transpose(Ad1))


x_est[:,2]=ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength,  x_model,u_model)
x_opl[:,2]=ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength,  x_model_op,u_model)

# measurements
y=np.zeros((2, Nsim))
y[0,:]=y_test[:-2,2]
y[1,:]=y_test[:-2,3]

C = np.zeros((2,number_of_states))
C[0,2] = 1
C[1,3] = 1


finish1 = time.perf_counter()
print('\n Finished offline in time',finish1 - start1,'seconds')

start2 = time.perf_counter()
#Nsim=398
for ti in range(2, Nsim-1):
    
    #############################update############################
    Lk=np.array((Pc[:,:,ti].dot(np.transpose(C))).dot(inv(C.dot(Pc[:,:,ti]).dot(np.transpose(C))+Rc)))
    x_est[:,ti-sequenceLength+1:ti+1]=x_est[:,ti-sequenceLength+1:ti+1] + Lk.dot((y[:,ti-1:ti+1]+v[ti,:,:]-C.dot(x_est[:,ti-sequenceLength+1:ti+1])))
    Pc[:,:,ti] = Pc[:,:,ti] - Lk.dot(C).dot(Pc[:,:,ti])
   
    u_model = X_test[ti,:,number_of_states:numberOfFeatures]
    x_model = x_est[:,ti-sequenceLength+1:ti+1].T
        
    Ad1= np.array(Ad_fun1(x_model[0],x_model[1],u_model))
    Ad2= np.array(Ad_fun2(x_model[0],x_model[1],u_model))
    Bd= np.array(Bd_fun(x_model[0],x_model[1],u_model))

    # ############################prediction########################
    Pc[:,:,ti+1] = Ad1.dot(Pc[:,:,ti-1]).dot(np.transpose(Ad1))+ Ad2.dot(Pc[:,:,ti]).dot(np.transpose(Ad2))+Q
    x_est[:,ti+1]=ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength,  x_model,u_model)
    x_opl[:,ti+1]=ObtainlstmModel_1(W_1, U_1, b_1, W_2, U_2, b_2, V_m, b_m, units_1, units_2, sequenceLength,  x_opl[:,ti-sequenceLength+1:ti+1].T,u_model)


finish2 = time.perf_counter()
print('\n Finished online in time',finish2 - start2,'seconds')

x_hat_unscaled = x_est.T*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]
x_opl_unscaled = x_opl.T*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]
x_act_unscaled = X_act.T*(data_max[0:number_of_states]-data_min[0:number_of_states]) + data_min[0:number_of_states]
x_test_unscaled = X_test*(data_max-data_min) + data_min


plt.figure(1)
plt.rcParams["figure.figsize"] = [5.00,3.00]
plt.rcParams["figure.autolayout"] = True
plt.plot(y_test_unscaled[0:Nsim,0],'r',label='Actual')
plt.plot(x_hat_unscaled[0:Nsim,0],'--g' ,label='Reduced_est',linewidth=2) 
plt.legend(loc='best')
plt.xlabel('Time step')
plt.ylabel("$X_{B3}$")
plt.show()


# The RMSE is reported for different initial guess and different random seed taking average
# This is a sample calculation

rmse_ekf=np.zeros(4)
for i in range(4):
    rmse_ekf[i]=rmse(x_hat_unscaled[:,i],y_test_unscaled[:-2,i])
print('rmse_entire_ekf',sum(rmse_ekf)/len(rmse_ekf))

rmse_cb_ekf=rmse(x_hat_unscaled[:,0],y_test_unscaled[:-2,0])
print('rmse_cb_ekf',rmse_cb_ekf)

# Degree of observability
O_A=control.obsv(Ad2,C)
rank_O=la.matrix_rank(O_A)
print('rank of A matrix=',rank_O)
print('conditioning no=' , np.linalg.cond(O_A))
print('Degree of observability=' , 1/np.linalg.cond(O_A))
finish = time.perf_counter()
print('\n Finished entire simulation in time',finish - start,'seconds')
   